import os
import timm
from utils.transforms import *
from  utils.yaml_tools import read_yaml, write_yaml
import os.path as osp


class CLASSFICATION_MODEL:
    # 这些参数可以写在配置文件中，通过配置文件的形式来进行查看和设置
    # 后期可以在初始化的时候加载这里的配置文件
    def __init__(self, model_name='efficientnet_b3', pretrained=False):
        # 数据集相关
        self.data_folder = "../CatDog"
        self.class_names = ['0', '1']
        self.num_classes = len(self.class_names)
        self.image_size = 224
        self.batch_size = 4
        data_transforms = get_torch_transforms(self.image_size)
        self.train_transforms = data_transforms['train']
        self.valid_transforms = data_transforms['val']
        # 模型训练相关参数
        self.model_name = model_name
        self.device = "cpu"
        self.model_save_path = "../checkpoints/"
        self.learning_rate = 0.0001
        self.epochs = 10
        self.weight_decay = 1e-5
        self.pretrained = pretrained
        # 初始化的时候任意建立一个模型占位
        self.model = timm.create_model(model_name, pretrained=False, num_classes=2)

    # 初始化加载的时候使用，目的是从本地读取配置文件并进行更新
    def update_model_info_from_local(self, config_path="configs/config.yaml"):
        # 每次调用这个函数的时候配置信息会进行一次重新刷新
        res_dict = read_yaml(config_path)
        # 1. 加载配置文件
        # 2. 将配置文件转换为对应的内容
        # 数据部分的更新
        self.data_folder = res_dict['data']['data_folder']
        self.class_names = os.listdir(os.path.join(self.data_folder, "train"))
        self.num_classes = len(self.class_names)
        self.batch_size = int(res_dict['data']['batch_size'])
        self.image_size = int(res_dict['data']['image_size'])
        # 模型部分更新
        self.model_name = res_dict['model']['model_name']
        self.pretrained = res_dict['model']['pretrained']

        self.model_save_path = res_dict['model']['model_save_path']
        os.makedirs(self.model_save_path, exist_ok=True)
        device_name = res_dict['model']['device']
        self.device = str(torch.device(device_name))
        # todo 判断gpu是否可以使用
        self.learning_rate = float(res_dict['model']['learning_rate'])
        self.epochs = int(res_dict['model']['epochs'])
        self.weight_decay = float(res_dict['model']['weight_decay'])
        # 初始化对应的模型，根据上面的信息。
        self.model = timm.create_model(self.model_name, pretrained=self.pretrained, num_classes=self.num_classes)
        # 3. 另外是关于配置文件的联动，配置文件联动的过程中需要结合
        # 其余的配置后面再补充，先将目前的配置搞定

    # 保存配置文件使用
    def save_local_config(self, config_path="configs/config.yaml"):
        # 将对应的信息保存在本地，下次加载的时候直接使用。
        print(type(self.batch_size))

        current_config = {'data': {'data_folder': self.data_folder,
                                    'image_size': self.image_size,
                                    'class_names': self.class_names,
                                    'num_classes': self.num_classes,
                                   'batch_size':self.batch_size},
                          'model': {'model_name': self.model_name,
                                    'device':self.device,
                                    'learning_rate': self.learning_rate,
                                    'epochs': self.epochs,
                                    'model_save_path': self.model_save_path,
                                    'weight_decay': self.weight_decay,
                                    'pretrained': self.pretrained},
                          'system': {'window_title': '分类系统'}}
        # 将当前的配置信息保存到本地
        write_yaml(current_config, config_path)
        # allow_unicode=True



    def data_init(self):
        if self.pretrained == True:
            save_dir = osp.join(self.model_save_path,
                                self.model_name + "-pretrained_" + str(self.image_size))  # 设置模型保存路径
        else:
            save_dir = osp.join(self.model_save_path,
                                self.model_name + "-nopretrained_" + str(self.image_size))  # 设置模型保存路径
        if not osp.isdir(save_dir):  # 如果保存路径不存在的话就创建
            os.makedirs(save_dir)  #
            print("save dir {} created".format(save_dir))

        # 加载数据集
        print(self.batch_size)
        print(type(self.batch_size))
        train_dataset = datasets.ImageFolder(osp.join(self.data_folder, "train"), self.train_transforms)  # 加载训练集
        valid_dataset = datasets.ImageFolder(osp.join(self.data_folder, "val"), self.valid_transforms)  # 加载验证集
        test_dataset = datasets.ImageFolder(osp.join(self.data_folder, "test"), self.valid_transforms)  # 加载测试集

        train_loader = DataLoader(  # 按照批次加载训练集
            train_dataset, batch_size=self.batch_size, shuffle=True,
            num_workers=0, pin_memory=True,
        )
        val_loader = DataLoader(  # 按照批次加载验证集
            valid_dataset, batch_size=self.batch_size, shuffle=False,
            num_workers=0, pin_memory=True,
        )

        test_loader = DataLoader(  # 按照批次加载验证集
            test_dataset, batch_size=self.batch_size, shuffle=False,
            num_workers=0, pin_memory=True,
        )
        return train_loader, val_loader, test_loader

    # def go_train(self):
    #     # 训练的时候占用数据集加载的内存，如果不训练，直接释放即可
    #     train_loader, val_loader, test_loader = self.data_init()
    #     # 由于训练过程是一个循环，并且训练需要记录状态，所以整个for循环还是需要写入到主的界面程序里面
    #     # 通过线程来启动，并且可以记录对应的状态










        # self.train_transforms = get_train_transforms()
# 模型训练


# 模型测试

# 模型推理